# Copyright (c) 2024 Huawei Technologies Co., Ltd.
#
# openMind is licensed under Mulan PSL v2.
# You can use this software according to the terms and conditions of the Mulan PSL v2.
# You may obtain a copy of Mulan PSL v2 at:
#
#          http://license.coscl.org.cn/MulanPSL2
#
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
# See the Mulan PSL v2 for more details.
import os

from openmind.utils.constants import Endpoint
from openmind.utils import logging

logger = logging.get_logger(__name__)


def get_platform_loader(target: str):

    openmind_platform = os.getenv("OPENMIND_PLATFORM")

    # openmind_hub will raise error when OPENMIND_PLATFORM is set to modelers, it should be None if download from modelers.
    if openmind_platform == Endpoint.modelers:
        raise ValueError("Please unset OPENMIND_PLATFORM if prepare downloading from  modelers.")

    # openmind_platform must have a value to determine correct repo in flow/model/loader.py and flow/datasets/loader.py
    if openmind_platform is None:
        openmind_platform = Endpoint.modelers

    if openmind_platform not in [Endpoint.modelers, Endpoint.hf]:
        raise ValueError(
            f"Unrecognized openmind platform {openmind_platform}. You can set OPENMIND_PLATFORM={Endpoint.hf} to download from huggingface , or unset OPENMIND_PLATFORM to download from modelers."
        )

    if target == "dataset":
        if openmind_platform == Endpoint.modelers:
            from openmind.integrations.datasets import load_dataset

            logger.info_rank0("Prepare downloading dataset from modelers")

            return load_dataset, openmind_platform
        elif openmind_platform == Endpoint.hf:
            from datasets import load_dataset

            logger.info_rank0("Prepare downloading dataset from huggingface")
            return load_dataset, openmind_platform

    elif target == "model":
        if openmind_platform == Endpoint.modelers:
            from openmind.utils.hub import snapshot_download

            logger.info_rank0("Prepare downloading model from modelers")

            return snapshot_download, openmind_platform
        elif openmind_platform == Endpoint.hf:
            from huggingface_hub import snapshot_download

            logger.info_rank0("Prepare downloading model from huggingface")

            return snapshot_download, openmind_platform
    else:
        raise ValueError("Currently supported target is model/dataset")
